import logging
import random
import math
import numpy as np
import matplotlib.pyplot as plt

from collections import defaultdict

from util.util_log import setup_logging
from util.util_csv import save_csv
from util.util_image import save_img, save_img_xyz, draw_img, plot_line_chart, save_line_chart
from util.util_ris_pattern import point_2_phi_pattern, phase_2_pattern, nRow, eps, phase_2_pattern_xyz, \
    phase_2_bit, phase_2_pattern_xyz_fft
from util.util_analysis_plane import get_peaks, get_peak_3rd, get_peak_5th, get_peak_9th

from multi_beam.multi_beam_NN import nn_beam_4, nn_beam_8_cub


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../../files/logs/log_multi_beam_NN.log")

# 获取日志记录器并记录日志
logger = logging.getLogger("[RIS-multi-beam-CS]")


# ============================================= 主函数的实现方法 =========================================
# NN位置
def find_positions(nn_template):
    """
    将二维数组nn_template中相同元素的位置放到同一列表并返回所有列表。
    参数:
    nn_template (np.array): 输入的二维数组。
    返回:
    list of list of int: 所有相同元素位置的列表集合。
    """
    positions = defaultdict(list)
    row_count, col_count = nn_template.shape
    # 遍历二维数组并记录元素的位置
    for i in range(row_count):
        for j in range(col_count):
            element = nn_template[i, j]
            position = i * col_count + j + 1  # 计算行优先顺序的位置
            positions[element].append(position)
    # 提取字典中的列表
    result = list(positions.values())
    return result


# 随机生成
def split_integers_randomly(k, n):
    # 生成包含1到k的整数列表
    integers = list(range(1, int(k) + 1))
    # 随机打乱整数列表
    np.random.shuffle(integers)
    # 初始化n个空数组
    result = [[] for _ in range(n)]
    # 将打乱后的整数逐一分配到n个数组中
    for idx, val in enumerate(integers):
        result[idx % n].append(val)
    return result


# 生成标记phase_mix的模板phase_mix_template
def init_phase_mix_template(M, N, k, m, n):
    # 创建一个空的MxN的np.array
    phase_mix_template = np.zeros((M, N), dtype=float)
    # 填充块的值
    block_num = 1
    for i in range(0, M, m):
        for j in range(0, N, n):
            # 避免超出块数量k
            if block_num > k:
                break
            phase_mix_template[i:i+m, j:j+n] = block_num
            block_num += 1
    return phase_mix_template


# 根据模板 phase_mix_template 和分类结果 clusters 生成 phase_mix
def generate_phase_mix(M, N, phase_mix_template, clusters):
    # 创建一个空的MxN的np.array用于存储结果
    phase_mix = np.zeros((M, N), dtype=float)
    # 创建一个字典，用于快速查找每个值属于哪个 phase
    cluster_dict = {}
    for cluster_info in clusters:
        phase = cluster_info['phase']
        cluster = cluster_info['cluster']
        for val in cluster:
            cluster_dict[val] = phase
    # 遍历 phase_mix_template，填充 phase_mix
    for i in range(M):
        for j in range(N):
            phase_id = int(phase_mix_template[i, j])
            if phase_id in cluster_dict:
                phase_mix[i, j] = cluster_dict[phase_id][i, j]
    return phase_mix


# ============================================= 遗传算法 ===========================================
class GeneticAlgorithm():
    __bit_num = 0       # 比特数
    __beam_num = 0      # 波束数
    __rows = 0
    __cols = 0
    __phase_mix_template = None
    __phases = None

    def __init__(self, bit_num, rows, cols, phase_mix_template, phases,
                 population_size=50, num_generations=100, num_parents=10, mutation_rate=0.1):
        # 初始化阵列相关
        self.__bit_num = bit_num
        self.__rows = rows
        self.__cols = cols
        self.__phase_mix_template = phase_mix_template
        self.__phases = phases
        self.__beam_num = len(phases)
        # 初始化遗传算法相关
        self.population_size = population_size
        self.num_generations = num_generations
        self.num_parents = num_parents
        self.mutation_rate = mutation_rate
        self.population = None
        self.best_individual = None
        self.best_fitness = None
        self.fitness_history = []  # 保存每一代的最佳适应度值

    def __get_psll_2(self, patternBit_mix):
        psll = 0
        # pattern 转dB
        pattern_dbw = 20 * np.log10(np.abs(patternBit_mix) / np.max(np.max(np.abs(patternBit_mix))) + eps)
        # pattern找峰值
        peaks = get_peaks(pattern_dbw)
        # 找第三峰(小于3dB)作为 PSLL
        if len(peaks) > 2:
            peak_3rd = get_peak_3rd(peaks)
            if peak_3rd is not None:
                psll = peak_3rd[0]
        return psll

    def __get_psll_4(self, patternBit_mix):
        psll = 0
        # pattern 转dB
        pattern_dbw = 20 * np.log10(np.abs(patternBit_mix) / np.max(np.max(np.abs(patternBit_mix))) + eps)
        # pattern找峰值
        peaks = get_peaks(pattern_dbw)
        # 找第三峰(小于3dB)作为 PSLL
        if len(peaks) > 4:
            peak_5th = get_peak_5th(peaks)
            if peak_5th is not None:
                psll = peak_5th[0]
        return psll

    def __get_psll_8(self, patternBit_mix):
        psll = 0
        # pattern 转dB
        pattern_dbw = 20 * np.log10(np.abs(patternBit_mix) / np.max(np.max(np.abs(patternBit_mix))) + eps)
        # pattern找峰值
        peaks = get_peaks(pattern_dbw)
        # 找第三峰(小于3dB)作为 PSLL
        if len(peaks) > 8:
            peak_9th = get_peak_9th(peaks)
            if peak_9th is not None:
                psll = peak_9th[0]
        return psll

    def __get_psll(self, patternBit_mix):
        psll = 0
        if self.__beam_num == 2:
            psll = self.__get_psll_2(patternBit_mix)
        elif self.__beam_num == 4:
            psll = self.__get_psll_4(patternBit_mix)
        elif self.__beam_num == 8:
            psll = self.__get_psll_8(patternBit_mix)
        return psll

    def fitness(self, cluster_init):
        # 生成 clusters
        clusters = []
        for i in range(len(self.__phases)):
            clusters.append({
                "phase": self.__phases[i],
                "cluster": cluster_init[i]
            })
        # 根据 分类结果cluster 和 模板phase_mix_template 生成 码阵phase_mix
        phase_mix = generate_phase_mix(self.__rows, self.__cols, self.__phase_mix_template, clusters)
        # 相位转换 X bit
        phaseBit_mix, phaseBitDeg_mix = phase_2_bit(phase_mix, self.__bit_num)
        # 计算phase_mix的方向图
        phaseBit_mix = np.deg2rad(phaseBitDeg_mix)
        # patternBit_mix = phase_2_pattern(phaseBit_mix)                # 公式法直接计算, 准确但速度太慢
        patternBit_mix, x, y = phase_2_pattern_xyz_fft(phaseBit_mix)    # FFT法计算, 快速
        # 计算码阵的最大副瓣
        psll = self.__get_psll(patternBit_mix)
        return psll

    # def fitness(self, cluster_init):
    #     list_sum = []
    #     for datas in cluster_init:
    #         sum = 0
    #         for data in datas:
    #             sum += data
    #         list_sum.append(sum)
    #     fit = 0
    #     for sum in list_sum:
    #         fit += sum ** 2
    #     return -fit

    def initialize_population(self, cluster_init):
        """初始化种群"""
        self.population = [cluster_init] * self.population_size

    def selection(self):
        """选择操作"""
        fitness_scores = [self.fitness(individual) for individual in self.population]
        sorted_indices = np.argsort(fitness_scores)  # 从低到高排序
        selected_parents = [self.population[i] for i in sorted_indices[:self.num_parents]]
        # 获取选中的父代个体对应的适应度值
        selected_fitness_scores = [fitness_scores[i] for i in sorted_indices[:self.num_parents]]
        return selected_parents, selected_fitness_scores

    def crossover(self, parents, offspring_size):
        """交叉操作"""
        offspring = []
        for _ in range(offspring_size):
            parent1, parent2 = random.sample(parents, 2)
            parent1 = np.array(parent1)
            parent2 = np.array(parent2)
            # 获取数组的形状
            shape = parent1.shape
            # 将 parent1 和 parent2 扁平化为一维数组
            flat_parent1 = parent1.flatten()
            flat_parent2 = parent2.flatten()
            # 随机选择一个切割位置
            cut_point = random.randint(0, len(flat_parent1))
            # 生成子代
            child = np.concatenate((flat_parent1[:cut_point], flat_parent2[cut_point:]))
            # 将子代重新塑形为原始的二维数组形状
            child = child.reshape(shape)
            child = child.tolist()
            #
            offspring.append(child)
        return offspring

    def mutation(self, offspring):
        """变异操作"""
        for individual in offspring:
            if random.random() < self.mutation_rate:
                # 以下执行变异操作
                # 获取数组的数量和长度
                num_clusters = len(individual)
                array_length = len(individual[0])
                for i in range(num_clusters):
                    for j in range(array_length):
                        # 随机选择另一个数组
                        other_cluster_index = random.choice([k for k in range(num_clusters) if k != i])
                        # 随机选择另一个数组中的一个位置
                        other_position = random.randint(0, array_length - 1)
                        # 交换元素
                        individual[i][j], individual[other_cluster_index][other_position] = \
                            individual[other_cluster_index][other_position], individual[i][j]
        return offspring

    def run(self, cluster_init):
        """运行遗传算法"""
        self.initialize_population(cluster_init)
        self.fitness_history = []
        self.best_fitness = self.fitness(cluster_init)
        self.best_individual = cluster_init
        logger.info("population_size=%d, num_generations=%d, num_parents=%d, mutation_rate=%d"
                    % (self.population_size, self.num_generations, self.num_parents, self.mutation_rate))
        for generation in range(self.num_generations):
            # 选择操作
            selected_parents, selected_fitness_scores = self.selection()
            # 交换操作
            offspring = self.crossover(selected_parents, self.population_size - self.num_parents)
            # 变异操作
            offspring = self.mutation(offspring)
            # 计算后代的适应度值
            offspring_fitness_scores = [self.fitness(individual) for individual in offspring]
            # 更新种群
            self.population = selected_parents + offspring
            # 合并父代和后代的适应度值
            all_fitness_scores = selected_fitness_scores + offspring_fitness_scores
            #
            # 找到当前代的最佳个体
            this_best_index = np.argmax(all_fitness_scores)  # 找到最佳适应度值的索引
            if this_best_index < len(selected_parents):
                this_best_individual = selected_parents[this_best_index]  # 最佳个体来自父代
            else:
                this_best_individual = offspring[this_best_index - len(selected_parents)]  # 最佳个体来自后代
            this_best_fitness = all_fitness_scores[this_best_index]  # 最佳适应度值
            #
            # 更新最佳适应度
            if this_best_fitness < self.best_fitness:
                self.best_fitness = this_best_fitness
                self.best_individual = this_best_individual
            # 记录最佳适应度曲线
            self.fitness_history.append(self.best_fitness)
            #
            logger.info("generation=%d: self.best_fitness=%f, self.best_individual:%s"
                        % (generation, self.best_fitness, self.best_individual))
        return self.best_individual, self.best_fitness, self.fitness_history


# ============================================= 主函数 ====================================
# NN 核心算法 -- 双波束
def cs_beam_2(phase1, phase2, bit_num):
    rows, cols = phase1.shape
    # 1.将phase_mix按平铺子阵划分方式分块, 模板为phase_mix_template
    m, n = 8, 8
    phase_mix_template = init_phase_mix_template(rows, cols, ((rows/m)*(cols/n)), m, n)
    # 2.生成填入方式
    # 随机填入
    # clusters_arr = split_integers_randomly(((rows/m)*(cols/n)), 2)
    # NN方式填入
    clusters_arr = []
    num = int((rows/m)*(cols/n))
    step = int(num // 2)
    clusters_arr.append(list(range(1, step + 1)))
    clusters_arr.append(list(range(step + 1, num + 1)))
    # 3.退火算法(SA)寻找最优的phase1和phase2的填入方式
    ga = GeneticAlgorithm(bit_num, rows, cols, phase_mix_template,
                          [phase1, phase2],
                          50, 500, 10, 0.1)
    best_individual, best_fitness, fitness_history = ga.run(clusters_arr)
    logger.info("best_individual:%s" % (best_individual))
    logger.info("Best Fitness=%f" % (best_fitness))
    # 4.填入phase1和phase2的块入phase_mix
    clusters = [{'phase': phase1, 'cluster': best_individual[0]},
                {'phase': phase2, 'cluster': best_individual[1]}]
    phase_mix = generate_phase_mix(rows, cols, phase_mix_template, clusters)
    #
    return phase_mix, phase_mix_template, best_individual, best_fitness, fitness_history


# NN 核心算法 -- 四波束
def cs_beam_4(phase1, phase2, phase3, phase4, bit_num):
    rows, cols = phase1.shape
    # 1.将phase_mix按平铺子阵划分方式分块, 模板为phase_mix_template
    m, n = 8, 8
    phase_mix_template = init_phase_mix_template(rows, cols, ((rows/m)*(cols/n)), m, n)
    # 2.生成填入方式
    # 随机填入
    # clusters_arr = split_integers_randomly(((rows/m)*(cols/n)), 4)
    # NN方式填入
    nn_template = nn_beam_4(np.full((rows//m, cols//n), 1), np.full((rows//m, cols//n), 2),
                            np.full((rows//m, cols//n), 3), np.full((rows//m, cols//n), 4))
    clusters_arr = find_positions(nn_template)
    # 3.退火算法(SA)寻找最优的phase1和phase2的填入方式
    ga = GeneticAlgorithm(bit_num, rows, cols, phase_mix_template,
                          [phase1, phase2, phase3, phase4],
                          50, 500, 10, 0.1)
    best_individual, best_fitness, fitness_history = ga.run(clusters_arr)
    logger.info("best_individual:%s" % (best_individual))
    logger.info("Best Fitness=%f" % (best_fitness))
    # 4.填入 phase1 到 phase4 的块入 phase_mix
    clusters = [{'phase': phase1, 'cluster': best_individual[0]},
                {'phase': phase2, 'cluster': best_individual[1]},
                {'phase': phase3, 'cluster': best_individual[2]},
                {'phase': phase4, 'cluster': best_individual[3]}]
    phase_mix = generate_phase_mix(rows, cols, phase_mix_template, clusters)
    #
    return phase_mix, phase_mix_template, best_individual, best_fitness, fitness_history


# NN 核心算法 -- 八波束
def cs_beam_8(phase1, phase2, phase3, phase4, phase5, phase6, phase7, phase8, bit_num):
    rows, cols = phase1.shape
    # 1.将phase_mix按平铺子阵划分方式分块, 模板为phase_mix_template
    m, n = 8, 8
    phase_mix_template = init_phase_mix_template(rows, cols, ((rows/m)*(cols/n)), m, n)
    # 2.生成填入方式
    # 随机填入
    # clusters_arr = split_integers_randomly(((rows/m)*(cols/n)), 8)
    # NN方式填入
    nn_template = nn_beam_8_cub(np.full((rows // m, cols // n), 1), np.full((rows // m, cols // n), 2),
                                np.full((rows // m, cols // n), 3), np.full((rows // m, cols // n), 4),
                                np.full((rows // m, cols // n), 5), np.full((rows // m, cols // n), 6),
                                np.full((rows // m, cols // n), 7), np.full((rows // m, cols // n), 8))
    clusters_arr = find_positions(nn_template)
    # 3.退火算法(SA)寻找最优的phase1和phase2的填入方式
    ga = GeneticAlgorithm(bit_num, rows, cols, phase_mix_template,
                          [phase1, phase2, phase3, phase4, phase5, phase6, phase7, phase8],
                          50, 500, 10, 0.1)
    best_individual, best_fitness, fitness_history = ga.run(clusters_arr)
    logger.info("best_individual:%s" % (best_individual))
    logger.info("Best Fitness=%f" % (best_fitness))
    # 4.填入 phase1 到 phase4 的块入 phase_mix
    clusters = [{'phase': phase1, 'cluster': best_individual[0]},
                {'phase': phase2, 'cluster': best_individual[1]},
                {'phase': phase3, 'cluster': best_individual[2]},
                {'phase': phase4, 'cluster': best_individual[3]},
                {'phase': phase5, 'cluster': best_individual[4]},
                {'phase': phase6, 'cluster': best_individual[5]},
                {'phase': phase7, 'cluster': best_individual[6]},
                {'phase': phase8, 'cluster': best_individual[7]}]
    phase_mix = generate_phase_mix(rows, cols, phase_mix_template, clusters)
    return phase_mix, phase_mix_template, best_individual, best_fitness, fitness_history


# 几何分区法 -- 双波束
def main_multi_beam_2(theta1, phi1, theta2, phi2, path_pre, bit_num):
    logger.info("main_multi_beam_2: bit_num=%d, path_pre=%s, " % (bit_num, path_pre))
    logger.info("main_multi_beam_2: theta1=%d, phi1=%d, theta2=%d, phi2=%d, " % (theta1, phi1, theta2, phi2))
    # 目前只支持2bit
    if bit_num > 2:
        logger.error("main_multi_beam_2: bit_num bigger than 2.")
        return
    phase1, phaseBit1, pattern1 = point_2_phi_pattern(theta1, phi1, bit_num)
    phase2, phaseBit2, pattern2 = point_2_phi_pattern(theta2, phi2, bit_num)
    # 确保 phase1 和 phase2 具有相同的形状
    assert phaseBit1.shape == phaseBit2.shape, "phase1 和 phase2 必须具有相同的形状"
    # NN-PS
    phase_mix, phase_mix_template, \
    best_cluster, best_fitness, fitness_history = cs_beam_2(phase1, phase2, bit_num)
    # 相位转换 X bit
    phaseBit_mix, phaseBitDeg_mix = phase_2_bit(phase_mix, bit_num)
    # 计算phase_mix的方向图
    phaseBit_mix = np.deg2rad(phaseBitDeg_mix)
    patternBit_mix = phase_2_pattern(phaseBit_mix)
    #
    # 保存结果
    logger.info("save NN multi-beam 2 result...")
    patternBit_mix_xyz, x, y, z = phase_2_pattern_xyz(phaseBit_mix)
    # 保存图片
    save_img(path_pre + "phase1.jpg", phase1)
    save_img(path_pre + "phase2.jpg", phase2)
    save_img(path_pre + "phaseBit1.jpg", phaseBit1)
    save_img(path_pre + "phaseBit2.jpg", phaseBit2)
    save_img(path_pre + "pattern1.jpg", pattern1)
    save_img(path_pre + "pattern2.jpg", pattern2)
    save_img(path_pre + "phaseBit_mix.jpg", phaseBit_mix)         # 几何分区法 -- 结果码阵
    save_img(path_pre + "patternBit_mix.jpg", patternBit_mix)     # 几何分区法 -- 结果码阵方向图
    save_img_xyz(path_pre + "patternBit_mix_xyz.jpg", np.abs(patternBit_mix_xyz), x, y)
    # 保存相位结果
    save_csv(phase1, path_pre + "phase1.csv")
    save_csv(phase2, path_pre + "phase2.csv")
    save_csv(phaseBit1, path_pre + "phaseBit1.csv")
    save_csv(phaseBit2, path_pre + "phaseBit2.csv")
    save_csv(phaseBit_mix, path_pre + "phaseBit_mix.csv")
    # 保存退火算法优化结果
    save_line_chart(path_pre + "fitness_history.jpg", fitness_history,
                    "fitness_history", "iteration", "best_fitness", "fitness")
    # 保存分区结果
    rows, cols = phaseBit1.shape
    phase_cub_1 = np.full((rows, cols), 0)
    phase_cub_2 = np.full((rows, cols), 180)
    clusters = [{'phase': phase_cub_1, 'cluster': best_cluster[0]},
                {'phase': phase_cub_2, 'cluster': best_cluster[1]}]
    phase_cub = generate_phase_mix(rows, cols, phase_mix_template, clusters)
    save_img(path_pre + "phase_cub.jpg", phase_cub)


# 几何分区法 -- 四波束
def main_multi_beam_4(theta1, phi1, theta2, phi2, theta3, phi3, theta4, phi4, path_pre, bit_num):
    logger.info("main_multi_beam_4: bit_num=%d, path_pre=%s, " % (bit_num, path_pre))
    logger.info("main_multi_beam_4: theta1=%d, phi1=%d, theta2=%d, phi2=%d, theta3=%d, phi3=%d, theta4=%d, phi4=%d"
                % (theta1, phi1, theta2, phi2, theta3, phi3, theta4, phi4))
    # 目前只支持2bit
    if bit_num > 2:
        logger.error("main_multi_beam_N: bit_num bigger than 2.")
        return
    phase1, phaseBit1, pattern1 = point_2_phi_pattern(theta1, phi1, bit_num)
    phase2, phaseBit2, pattern2 = point_2_phi_pattern(theta2, phi2, bit_num)
    phase3, phaseBit3, pattern3 = point_2_phi_pattern(theta3, phi3, bit_num)
    phase4, phaseBit4, pattern4 = point_2_phi_pattern(theta4, phi4, bit_num)
    # 确保所有数组具有相同的形状
    assert phaseBit1.shape == phaseBit2.shape == phaseBit3.shape == phaseBit4.shape, "所有数组必须具有相同的形状"
    # NN - S - SA
    phase_mix, phase_mix_template, best_cluster, \
    best_fitness, fitness_history = cs_beam_4(phase1, phase2, phase3, phase4, bit_num)
    # 相位转换 X bit
    phaseBit_mix, phaseBitDeg_mix = phase_2_bit(phase_mix, bit_num)
    # 计算phase_mix的方向图
    phaseBit_mix = np.deg2rad(phaseBitDeg_mix)
    # 计算phase_mix
    patternBit_mix = phase_2_pattern(phaseBit_mix)
    #
    # 保存结果
    logger.info("save NN multi-beam 4 result...")
    patternBit_mix_xyz, x, y, z = phase_2_pattern_xyz(phaseBit_mix)
    # 保存图片
    save_img(path_pre + "phase1.jpg", phase1)
    save_img(path_pre + "phase2.jpg", phase2)
    save_img(path_pre + "phase3.jpg", phase3)
    save_img(path_pre + "phase4.jpg", phase4)
    save_img(path_pre + "phaseBit1.jpg", phaseBit1)
    save_img(path_pre + "phaseBit2.jpg", phaseBit2)
    save_img(path_pre + "phaseBit3.jpg", phaseBit3)
    save_img(path_pre + "phaseBit4.jpg", phaseBit4)
    save_img(path_pre + "pattern1.jpg", pattern1)
    save_img(path_pre + "pattern2.jpg", pattern2)
    save_img(path_pre + "pattern3.jpg", pattern3)
    save_img(path_pre + "pattern4.jpg", pattern4)
    save_img(path_pre + "phase_mix.jpg", phaseBit_mix)       # 几何分区法 -- 结果码阵
    save_img(path_pre + "pattern_mix.jpg", patternBit_mix)   # 几何分区法 -- 结果码阵方向图
    save_img_xyz(path_pre + "patternBit_mix_xyz.jpg", np.abs(patternBit_mix_xyz), x, y)
    # 保存相位结果
    save_csv(phase1, path_pre + "phase1.csv")
    save_csv(phase2, path_pre + "phase2.csv")
    save_csv(phase3, path_pre + "phase3.csv")
    save_csv(phase4, path_pre + "phase4.csv")
    save_csv(phaseBit1, path_pre + "phaseBit1.csv")
    save_csv(phaseBit2, path_pre + "phaseBit2.csv")
    save_csv(phaseBit3, path_pre + "phaseBit3.csv")
    save_csv(phaseBit4, path_pre + "phaseBit4.csv")
    save_csv(phaseBit_mix, path_pre + "phase_mix.csv")
    # 保存退火算法优化结果
    save_line_chart(path_pre + "fitness_history.jpg", fitness_history,
                    "fitness_history", "iteration", "best_fitness", "fitness")
    # 保存分区结果
    rows, cols = phaseBit1.shape
    phase_cub_1 = np.full((rows, cols), 0)
    phase_cub_2 = np.full((rows, cols), 90)
    phase_cub_3 = np.full((rows, cols), 180)
    phase_cub_4 = np.full((rows, cols), 270)
    clusters = [{'phase': phase_cub_1, 'cluster': best_cluster[0]},
                {'phase': phase_cub_2, 'cluster': best_cluster[1]},
                {'phase': phase_cub_3, 'cluster': best_cluster[2]},
                {'phase': phase_cub_4, 'cluster': best_cluster[3]}]
    phase_cub = generate_phase_mix(rows, cols, phase_mix_template, clusters)
    save_img(path_pre + "phase_cub.jpg", phase_cub)


def main_multi_beam_8(theta1, phi1, theta2, phi2, theta3, phi3, theta4, phi4,
                      theta5, phi5, theta6, phi6, theta7, phi7, theta8, phi8,
                      path_pre, bit_num):
    logger.info("main_multi_beam_8: bit_num=%d, path_pre=%s, " % (bit_num, path_pre))
    logger.info("main_multi_beam_8: theta1=%d, phi1=%d, theta2=%d, phi2=%d, theta3=%d, phi3=%d, theta4=%d, phi4=%d, "
                "theta5=%d, phi5=%d, theta6=%d, phi6=%d, theta7=%d, phi7=%d, theta8=%d, phi8=%d"
                % (theta1, phi1, theta2, phi2, theta3, phi3, theta4, phi4,
                   theta5, phi5, theta6, phi6, theta7, phi7, theta8, phi8))
    # 目前只支持2bit
    if bit_num > 2:
        logger.error("main_multi_beam_8: bit_num bigger than 2.")
        return
    # 获取所有的 phaseBit 变量
    phase1, phaseBit1, pattern1 = point_2_phi_pattern(theta1, phi1, bit_num)
    phase2, phaseBit2, pattern2 = point_2_phi_pattern(theta2, phi2, bit_num)
    phase3, phaseBit3, pattern3 = point_2_phi_pattern(theta3, phi3, bit_num)
    phase4, phaseBit4, pattern4 = point_2_phi_pattern(theta4, phi4, bit_num)
    phase5, phaseBit5, pattern5 = point_2_phi_pattern(theta5, phi5, bit_num)
    phase6, phaseBit6, pattern6 = point_2_phi_pattern(theta6, phi6, bit_num)
    phase7, phaseBit7, pattern7 = point_2_phi_pattern(theta7, phi7, bit_num)
    phase8, phaseBit8, pattern8 = point_2_phi_pattern(theta8, phi8, bit_num)
    # 确保所有数组具有相同的形状
    assert phaseBit1.shape == phaseBit2.shape == phaseBit3.shape == phaseBit4.shape == \
           phaseBit5.shape == phaseBit6.shape == phaseBit7.shape == phaseBit8.shape, "所有数组必须具有相同的形状"
    # NN - S - SA
    phase_mix, phase_mix_template, best_cluster, best_fitness, fitness_history \
        = cs_beam_8(phase1, phase2, phase3, phase4, phase5, phase6, phase7, phase8, bit_num)
    # 相位转换 X bit
    phaseBit_mix, phaseBitDeg_mix = phase_2_bit(phase_mix, bit_num)
    # 计算phase_mix的方向图
    phaseBit_mix = np.deg2rad(phaseBitDeg_mix)
    # 计算phase_mix
    patternBit_mix = phase_2_pattern(phaseBit_mix)
    # 保存结果
    logger.info("save NN multi-beam 8 result...")
    patternBit_mix_xyz, x, y, z = phase_2_pattern_xyz(phaseBit_mix)
    # 保存图片
    save_img(path_pre + "phase1.jpg", phase1)
    save_img(path_pre + "phase2.jpg", phase2)
    save_img(path_pre + "phase3.jpg", phase3)
    save_img(path_pre + "phase4.jpg", phase4)
    save_img(path_pre + "phase5.jpg", phase5)
    save_img(path_pre + "phase6.jpg", phase6)
    save_img(path_pre + "phase7.jpg", phase7)
    save_img(path_pre + "phase8.jpg", phase8)
    save_img(path_pre + "phaseBit1.jpg", phaseBit1)
    save_img(path_pre + "phaseBit2.jpg", phaseBit2)
    save_img(path_pre + "phaseBit3.jpg", phaseBit3)
    save_img(path_pre + "phaseBit4.jpg", phaseBit4)
    save_img(path_pre + "phaseBit5.jpg", phaseBit5)
    save_img(path_pre + "phaseBit6.jpg", phaseBit6)
    save_img(path_pre + "phaseBit7.jpg", phaseBit7)
    save_img(path_pre + "phaseBit8.jpg", phaseBit8)
    save_img(path_pre + "pattern1.jpg", pattern1)
    save_img(path_pre + "pattern2.jpg", pattern2)
    save_img(path_pre + "pattern3.jpg", pattern3)
    save_img(path_pre + "pattern4.jpg", pattern4)
    save_img(path_pre + "pattern5.jpg", pattern5)
    save_img(path_pre + "pattern6.jpg", pattern6)
    save_img(path_pre + "pattern7.jpg", pattern7)
    save_img(path_pre + "pattern8.jpg", pattern8)
    save_img(path_pre + "phaseBit_mix.jpg", phaseBit_mix)  # 几何分区法 -- 结果码阵
    save_img(path_pre + "patternBit_mix.jpg", patternBit_mix)  # 几何分区法 -- 结果码阵方向图
    save_img_xyz(path_pre + "patternBit_mix_xyz.jpg", np.abs(patternBit_mix_xyz), x, y)
    # 保存相位结果
    save_csv(phase1, path_pre + "phase1.csv")
    save_csv(phase2, path_pre + "phase2.csv")
    save_csv(phase3, path_pre + "phase3.csv")
    save_csv(phase4, path_pre + "phase4.csv")
    save_csv(phase5, path_pre + "phase5.csv")
    save_csv(phase6, path_pre + "phase6.csv")
    save_csv(phase7, path_pre + "phase7.csv")
    save_csv(phase8, path_pre + "phase8.csv")
    save_csv(phaseBit1, path_pre + "phaseBit1.csv")
    save_csv(phaseBit2, path_pre + "phaseBit2.csv")
    save_csv(phaseBit3, path_pre + "phaseBit3.csv")
    save_csv(phaseBit4, path_pre + "phaseBit4.csv")
    save_csv(phaseBit5, path_pre + "phaseBit5.csv")
    save_csv(phaseBit6, path_pre + "phaseBit6.csv")
    save_csv(phaseBit7, path_pre + "phaseBit7.csv")
    save_csv(phaseBit8, path_pre + "phaseBit8.csv")
    save_csv(phaseBit_mix, path_pre + "phaseBit_mix.csv")
    # 保存退火算法优化结果
    save_line_chart(path_pre + "fitness_history.jpg", fitness_history,
                    "fitness_history", "iteration", "best_fitness", "fitness")
    # 保存分区结果
    rows, cols = phaseBit1.shape
    phase_cub_1 = np.full((rows, cols), 0)
    phase_cub_2 = np.full((rows, cols), 45)
    phase_cub_3 = np.full((rows, cols), 90)
    phase_cub_4 = np.full((rows, cols), 135)
    phase_cub_5 = np.full((rows, cols), 180)
    phase_cub_6 = np.full((rows, cols), 225)
    phase_cub_7 = np.full((rows, cols), 270)
    phase_cub_8 = np.full((rows, cols), 315)
    clusters = [{'phase': phase_cub_1, 'cluster': best_cluster[0]},
                {'phase': phase_cub_2, 'cluster': best_cluster[1]},
                {'phase': phase_cub_3, 'cluster': best_cluster[2]},
                {'phase': phase_cub_4, 'cluster': best_cluster[3]},
                {'phase': phase_cub_5, 'cluster': best_cluster[4]},
                {'phase': phase_cub_6, 'cluster': best_cluster[5]},
                {'phase': phase_cub_7, 'cluster': best_cluster[6]},
                {'phase': phase_cub_8, 'cluster': best_cluster[7]}]
    phase_cub = generate_phase_mix(rows, cols, phase_mix_template, clusters)
    save_img(path_pre + "phase_cub.jpg", phase_cub)



# ============================================= 测试函数 ====================================
def test_init_phase_mix():
    # 1.将phase_mix按平铺子阵划分方式分块
    M, N = 64, 64
    m, n = 8, 8
    phase_mix_template = init_phase_mix_template(M, N, ((M/m)*(N/n)), m, n)
    phase1 = np.full((M, N), 22)
    phase2 = np.full((M, N), 44)
    clusters_arr = []
    rows, cols = 64, 64
    num = int((rows / m) * (cols / n))
    step = int(num // 2)
    clusters_arr.append(list(range(1, step + 1)))
    clusters_arr.append(list(range(step + 1, num + 1)))
    clusters = [{'phase': phase1, 'cluster': clusters_arr[0]},
                {'phase': phase2, 'cluster': clusters_arr[1]}]
    phase_mix = generate_phase_mix(M, N, phase_mix_template, clusters)
    draw_img(phase_mix_template)
    draw_img(phase_mix)
    #
    # M, N = 64, 64
    # m, n = 32, 16
    # phase_mix_template = init_phase_mix_template(M, N, ((M / m) * (N / n)), m, n)
    # phase1 = np.full((M, N), 22)
    # phase2 = np.full((M, N), 44)
    # phase3 = np.full((M, N), 66)
    # clusters = [{'phase': phase1, 'cluster': [1, 2, 3]},
    #             {'phase': phase2, 'cluster': [4, 5, 6]},
    #             {'phase': phase3, 'cluster': [7, 8]}]
    # phase_mix = generate_phase_mix(M, N, phase_mix_template, clusters)
    # draw_img(phase_mix_template)
    # draw_img(phase_mix)


def test_split_integers_randomly():
    # 打印结果的辅助函数
    def print_split_results(k, n, result):
        for i in range(n):
            print(f"arr{i + 1}={result[i]}")
    # 示例1
    k, n = 8, 2
    result = split_integers_randomly(k, n)
    print(f"k={k}, n={n}")
    print_split_results(k, n, result)
    print("\n")
    # 示例2
    k, n = 8, 4
    result = split_integers_randomly(k, n)
    print(f"k={k}, n={n}")
    print_split_results(k, n, result)
    print("\n")
    # 示例3
    k, n = 6, 2
    result = split_integers_randomly(k, n)
    print(f"k={k}, n={n}")
    print_split_results(k, n, result)
    print("\n")
    # 示例4
    k, n = 6, 3
    result = split_integers_randomly(k, n)
    print(f"k={k}, n={n}")
    print_split_results(k, n, result)
    print("\n")


def test_cs_beam_2():
    M, N = 64, 64
    phase1 = np.full((M, N), 22)
    phase2 = np.full((M, N), 44)
    phase_mix, phase_mix_template, best_cluster, best_fitness, fitness_history \
        = cs_beam_2(phase1, phase2, 1)
    draw_img(phase_mix)
    plot_line_chart(fitness_history, "fitness_history", "iterations", "fitness", "best_fitness")


def test_cs_beam_4():
    M, N = 64, 64
    phase1 = np.full((M, N), 22)
    phase2 = np.full((M, N), 44)
    phase3 = np.full((M, N), 66)
    phase4 = np.full((M, N), 88)
    phase_mix, phase_mix_template, clusters_arr, best_fitness, fitness_history \
        = cs_beam_4(phase1, phase2, phase3, phase4, 1)
    draw_img(phase_mix)


def test_cs_beam_8():
    M, N = 64, 64
    phase1 = np.full((M, N), 11)
    phase2 = np.full((M, N), 22)
    phase3 = np.full((M, N), 33)
    phase4 = np.full((M, N), 44)
    phase5 = np.full((M, N), 55)
    phase6 = np.full((M, N), 66)
    phase7 = np.full((M, N), 77)
    phase8 = np.full((M, N), 88)
    phase_mix, phase_mix_template, clusters_arr, best_fitness, fitness_history \
        = cs_beam_8(phase1, phase2, phase3, phase4, phase5, phase6, phase7, phase8, 1)
    draw_img(phase_mix)


def test_GA():
    bit_num = 1
    phase1, phaseBit1, pattern1 = point_2_phi_pattern(30, 0, bit_num)
    phase2, phaseBit2, pattern2 = point_2_phi_pattern(30, 180, bit_num)
    rows, cols = phase1.shape
    m, n = 8, 8
    phase_mix_template = init_phase_mix_template(rows, cols, ((rows / m) * (cols / n)), m, n)
    cluster_init = []
    cluster_num = 2
    num = int((rows/m)*(cols/n))
    step = int(num // cluster_num)
    cluster_init.append(list(range(1, step + 1)))
    cluster_init.append(list(range(step + 1, num + 1)))
    ga = GeneticAlgorithm(bit_num, rows, cols, phase_mix_template, [phase1, phase2], 50, 100, 10, 0.1)
    best_individual, best_fitness, fitness_history = ga.run(cluster_init)
    logger.info("best_individual:%s" % (best_individual))
    logger.info("Best Fitness=%f" % (best_fitness))
    plot_line_chart(fitness_history, "fitness_history", "iteration", "best_fitness", "fitness")


# ======================================================= main 主方法 ===============================================
def main_multi_nn_s_ga():
    # 基于NN-S-GA的方法: 主函数
    # 测试方法
    # test_init_phase_mix()
    # test_split_integers_randomly()
    # test_cs_beam_2()
    # test_cs_beam_4()
    # test_cs_beam_8()
    # test_GA()
    #
    # 几何分区法: 主函数
    main_multi_beam_2(30, 0, 30, 90,
                      "../files/multi-beam/1bit/NN-S-GA/2-(30,0,30,90)/", 1)
    main_multi_beam_2(30, 0, 30, 180,
                      "../files/multi-beam/1bit/NN-S-GA/2-(30,0,30,180)/", 1)
    main_multi_beam_4(30, 0, 30, 60, 30, 120, 30, 180,
                      "../files/multi-beam/1bit/NN-S-GA/4-(30,0,30,60,30,120,30,180)/", 1)
    main_multi_beam_4(30, 0, 30, 90, 30, 180, 30, 270,
                      "../files/multi-beam/1bit/NN-S-GA/4-(30,0,30,90,30,180,30,270)/", 1)
    main_multi_beam_8(30, 0, 30, 45, 30, 90, 30, 135, 30, 180, 30, 225, 30, 270, 30, 315,
                      "../files/multi-beam/1bit/NN-S-GA/8-(30,45step)/", 1)
    #
    main_multi_beam_2(30, 0, 30, 90,
                      "../files/multi-beam/2bit/NN-S-GA/2-(30,0,30,90)/", 2)
    main_multi_beam_2(30, 0, 30, 180,
                      "../files/multi-beam/2bit/NN-S-GA/2-(30,0,30,180)/", 2)
    main_multi_beam_4(30, 0, 30, 60, 30, 120, 30, 180,
                      "../files/multi-beam/2bit/NN-S-GA/4-(30,0,30,60,30,120,30,180)/", 2)
    main_multi_beam_4(30, 0, 30, 90, 30, 180, 30, 270,
                      "../files/multi-beam/2bit/NN-S-GA/4-(30,0,30,90,30,180,30,270)/", 2)
    main_multi_beam_8(30, 0, 30, 45, 30, 90, 30, 135, 30, 180, 30, 225, 30, 270, 30, 315,
                      "../files/multi-beam/2bit/NN-S-GA/8-(30,45step)/", 2)




if __name__ == '__main__':
    logger.info("1bit-RIS-multi-beam-CS: Cluster")
    main_multi_nn_s_ga()